package com.maiko.maikoaiagent.rag;

import jakarta.annotation.Resource;
import org.springframework.ai.document.Document;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.pgvector.PgVectorStore;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.jdbc.core.JdbcTemplate;

import java.util.List;

import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgDistanceType.COSINE_DISTANCE;
import static org.springframework.ai.vectorstore.pgvector.PgVectorStore.PgIndexType.HNSW;

/**
 * 配置类，用于创建和配置 PgVectorStore 实例
 * @author: Maiko7
 * @create: 2025-06-16 7:45
 */
@Configuration
public class PgVectorVectorStoreConfig {

    @Resource
    private LoveAppDocumentLoader loveAppDocumentLoader;

    /**
     * 创建并配置一个 VectorStore Bean
     * 使用 PgVectorStore 作为向量存储实现，适用于 PostgreSQL 数据库
     *
     * @param jdbcTemplate            Spring 提供的 JDBC 模板，用于数据库操作
     * @param dashscopeEmbeddingModel 嵌入模型，用于将数据转换为向量表示
     * @return 返回配置好的 VectorStore 实例
     */
    @Bean
    public VectorStore pgVectorVectorStore(JdbcTemplate jdbcTemplate, EmbeddingModel dashscopeEmbeddingModel) {
        // 构建 PgVectorStore 实例，并设置相关配置参数
        VectorStore vectorStore = PgVectorStore.builder(jdbcTemplate, dashscopeEmbeddingModel)
                .dimensions(1536)                    // 设置向量维度，默认为 1536
                .distanceType(COSINE_DISTANCE)       // 设置距离计算方式为余弦相似度
                .indexType(HNSW)                     // 设置索引类型为 HNSW（高效近似最近邻搜索）
                .initializeSchema(true)              // 是否初始化数据库表结构（首次运行时建议开启）
                .schemaName("public")                // 设置数据库 schema 名称，默认为 "public"
                .vectorTableName("vector_store")     // 设置向量存储表名，默认为 "vector_store"
                .maxDocumentBatchSize(10000)         // 设置批量处理文档的最大数量，默认为 10000
                .build();

        // 加载文档
        List<Document> documents = loveAppDocumentLoader.loadMarkdowns();
        vectorStore.add(documents);
        return vectorStore;
    }
}


